import sys


sys.path.append("/data1/yxn/repast-city/")
from typing import Dict, Callable, Type
from envs.meituan import MeituanRiderWrapper
from casual.envs.meituan_micro import MeituanMicro
from casual.envs.meituan_macro import MeituanMacro
from casual.envs.meituan_reposition import MeituanRepo
from envs.lunar_lander import LunarLander
from envs.sc2_biuld_marines import SC2BuildMarine
from envs.cartpole import Cartpole
from envs.aimtest import AimTestEnv
from core import Env
from learning.config import Config



ALL_ENVS: Dict[str, Type[Env]] = {
    'lunarlander': LunarLander,
    'cartpole': Cartpole,
    'buildmarine': SC2BuildMarine,
    'aimtest': AimTestEnv,
    'meituan_micro': MeituanMicro,
    'meituan_macro': MeituanMacro,
    'meituan2': MeituanRiderWrapper,
    'meituan_repo': MeituanRepo
}


def get_env_class(env_id: str):
    try:
        return ALL_ENVS[env_id]
    except KeyError as e:
        print(f"unsupported environment id '{env_id}'")
        print("supported environments are: ", sep='')
        print(', '.join(ALL_ENVS.keys()))
        raise e


def get_default_config(env_name: str, args):
    
    config = Config()
    config.device_id = 'cuda'
    config.rl.n_epoch_actor = 2 # Actor 网络的训练轮数
    config.rl.n_epoch_critic = 16 # Critic 网络的训练轮数
    config.rl.optim.lr = 1e-4 # 强化学习优化器的学习率
    config.model.pthres = 0.2 # 因果发现确认父节点的阈值
    # config.model.pthres_max = 0.5
    # config.model.pthres_min = 0.15
    config.rl.optim.batchsize = 512 # 强化学习优化器的批处理大小，即每次更新模型时使用的样本数量。
    config.model.optim.lr = 1e-4 # 模型优化器的学习率
    config.model.optim.max_grad_norm = 1.0 # 最大梯度范数，用于梯度裁剪，防止梯度爆炸。
    config.mbrl.explore_rate_max = 0.8 # 最大探索率，控制模型在训练过程中探索新策略的比例。
    config.mbrl.n_round_planning = 20 # 规划的轮数，决定模型在每个训练步骤中规划的次数。
    config.mbrl.causal_interval_min = 3 # 最小因果间隔，用于因果关系建模
    config.mbrl.causal_interval_max = 3 # 最大因果间隔。
    config.mbrl.causal_interval_increase = 0.1 # 因果间隔的增加量
    config.mbrl.ensemble_size = 5 # 模型集合的大小，表示同时训练的多个模型的数量。
    config.model.n_jobs_fcit = 4 # 模型中用于因果关系测试的并行作业数。

    if env_name == 'cartpole':
        config.dims.variable_encoder_hidden = 64 # 变量编码器的隐藏层大小。
        config.dims.variable_encoding = 64 # 变量编码的大小。
        config.rl.n_sample = 1024 # 每次训练中使用的样本数量。
        config.rl.discount = 0.98 # 折扣因子，用于计算未来奖励的现值
        config.rl.gae_lambda = 0.95 # 广义优势估计中的 λ 值。
        config.rl.kl_penalty = 0.2 # KL 散度惩罚项，用于控制策略更新的幅度。
        config.rl.entropy_penalty = 0.04 # 熵惩罚项，用于增加策略的探索性。
        config.rl.max_episode_length = 200 # 每个 episode 的最大长度。
        config.model.buffer_size = 200000 # 经验缓冲区的大小，用于存储训练数据。
        config.model.optim.batchsize = 1024 # 模型优化器的批处理大小。
        # config.model.n_sample_oracle = 5000
        config.mbrl.n_batch_fit =  400 # 拟合模型所需的批次数。
        config.mbrl.n_batch_fit_new_graph = 800 # 拟合新图所需的批次数。
        config.mbrl.n_sample_explore = 400 # 探索过程中使用的样本数量。
        config.mbrl.n_sample_exploit = 400 # 开发过程中使用的样本数量。
        config.mbrl.n_sample_warmup = 1200 # 热身阶段使用的样本数量。
        config.mbrl.n_sample_rollout = 4096 # rollout 过程中使用的样本数量。
        config.mbrl.rollout_length = (1, 5) # rollout 的长度范围。
    elif env_name == 'buildmarine':
        config.dims.variable_encoder_hidden = 64
        config.dims.variable_encoding = 64
        config.rl.n_sample = 128
        config.rl.discount = 0.95
        config.rl.gae_lambda = 0.9
        config.rl.kl_penalty = 0.2
        config.rl.entropy_penalty = 0.04
        config.rl.max_episode_length = 40
        config.model.buffer_size = 100000
        config.model.optim.batchsize = 1024
        # config.model.n_sample_oracle = 10000
        config.mbrl.n_batch_fit = 400
        config.mbrl.n_batch_fit_new_graph = 800
        config.mbrl.n_sample_explore = 64
        config.mbrl.n_sample_exploit = 64
        config.mbrl.n_sample_warmup = 512
        config.mbrl.n_sample_rollout = 2048
        config.mbrl.rollout_length = (3, 8)
    elif env_name == 'lunarlander':
        config.dims.variable_encoder_hidden = 128
        config.dims.variable_encoding = 128
        config.rl.n_sample = 2048
        config.rl.discount = 0.975
        config.rl.gae_lambda = 0.94
        config.rl.kl_penalty = 0.2
        config.rl.entropy_penalty = 0.005 if args.continuous else 0.04
        config.rl.max_episode_length = 512
        config.model.buffer_size = 200000
        config.model.optim.batchsize = 1024
        # config.model.n_sample_oracle = 200000
        config.mbrl.n_batch_fit =  400
        config.mbrl.n_batch_fit_new_graph = 800
        config.mbrl.n_sample_explore = 1024
        config.mbrl.n_sample_exploit = 1024
        config.mbrl.n_sample_warmup = 4096
        config.mbrl.n_sample_rollout = 4096
        config.mbrl.rollout_length = (1, 10)
    elif env_name == 'aimtest':
        config.dims.variable_encoder_hidden = 64
        config.dims.variable_encoding = 64
        config.dims.inferrer_key = 64
        config.dims.inferrer_value = 64
        config.dims.distribution_embedding = 32
        config.dims.decoder_hidden =32
        config.model.buffer_size = 200000
        config.model.optim.batchsize = 1024
        config.model.pthres = 0.2
        # config.model.n_sample_oracle = 200000
    elif env_name == 'meituan_micro':
        config.dims.variable_encoder_hidden = 128 
        config.dims.variable_encoding = 128 
        config.rl.n_sample = 360 * 5 # 样本数量，以确保模型在每个 epoch 内有足够的数据训练。
        config.rl.discount = 0.975 
        config.rl.gae_lambda = 0.95 
        config.rl.kl_penalty = 0.2 
        config.rl.entropy_penalty = 0.01 # 增加熵惩罚，以促进探索。
        config.rl.max_episode_length = 360 * 5 # episode 的最大长度，为3天时间
        config.model.buffer_size = 200000 
        config.model.optim.batchsize = 180 # 减少模型优化器的批处理大小，以增加更新频率。
        # config.model.n_sample_oracle = 200000
        config.mbrl.n_batch_fit = 5 # 增加拟合批次数，以确保模型训练的充分性。
        config.mbrl.n_batch_fit_new_graph = 10 # 增加拟合新图的批次数。
        config.mbrl.n_sample_explore = 360 * 3 # 增加探索过程中使用的样本数量。
        config.mbrl.n_sample_exploit = 360 * 3 # 增加开发过程中使用的样本数量。
        config.mbrl.n_sample_warmup = 360 * 3 # 增加热身阶段使用的样本数量。
        config.mbrl.n_sample_rollout = 360 * 3# rollout 过程中使用的样本数量。
        config.mbrl.rollout_length = (1, 10) # 调整 rollout 长度，以平衡样本数量和训练时间。
    elif env_name == 'meituan_macro':
        config.dims.variable_encoder_hidden = 128 
        config.dims.variable_encoding = 128 
        config.rl.n_sample = 360 * 3 # 样本数量，以确保模型在每个 epoch 内有足够的数据训练。
        config.rl.discount = 0.975 
        config.rl.gae_lambda = 0.95 
        config.rl.kl_penalty = 0.2 
        config.rl.entropy_penalty = 0.01 # 增加熵惩罚，以促进探索。
        config.rl.max_episode_length = 360 * 3 # episode 的最大长度，为3天时间
        config.model.buffer_size = 200000 
        config.model.optim.batchsize = 180 # 减少模型优化器的批处理大小，以增加更新频率。
        # config.model.n_sample_oracle = 200000
        config.mbrl.n_batch_fit = 5 # 增加拟合批次数，以确保模型训练的充分性。
        config.mbrl.n_batch_fit_new_graph = 10 # 增加拟合新图的批次数。
        config.mbrl.n_sample_explore = 360 # 增加探索过程中使用的样本数量。
        config.mbrl.n_sample_exploit = 360 # 增加开发过程中使用的样本数量。
        config.mbrl.n_sample_warmup = 360 # 增加热身阶段使用的样本数量。
        config.mbrl.n_sample_rollout = 360 # rollout 过程中使用的样本数量。
        config.mbrl.rollout_length = (1, 10) # 调整 rollout 长度，以平衡样本数量和训练时间。
    elif env_name == 'meituan_repo':
        config.dims.variable_encoder_hidden = 128 
        config.dims.variable_encoding = 128 
        config.rl.n_sample = 360 * 3 # 样本数量，以确保模型在每个 epoch 内有足够的数据训练。
        config.rl.discount = 0.975 
        config.rl.gae_lambda = 0.95 
        config.rl.kl_penalty = 0.2 
        config.rl.entropy_penalty = 0.01 # 增加熵惩罚，以促进探索。
        config.rl.max_episode_length = 360 * 3 # episode 的最大长度，为3天时间
        config.model.buffer_size = 200000 
        config.model.optim.batchsize = 180 # 减少模型优化器的批处理大小，以增加更新频率。
        # config.model.n_sample_oracle = 200000
        config.mbrl.n_batch_fit = 5 # 增加拟合批次数，以确保模型训练的充分性。
        config.mbrl.n_batch_fit_new_graph = 10 # 增加拟合新图的批次数。
        config.mbrl.n_sample_explore = 360  # 增加探索过程中使用的样本数量。
        config.mbrl.n_sample_exploit = 360  # 增加开发过程中使用的样本数量。
        config.mbrl.n_sample_warmup = 360 # 增加热身阶段使用的样本数量。
        config.mbrl.n_sample_rollout = 360 # rollout 过程中使用的样本数量。
        config.mbrl.rollout_length = (1, 10) # 调整 rollout 长度，以平衡样本数量和训练时间。
        config.rl.kl_penalty = 0.6
    return config
